import cv2
import numpy as np

# 打开摄像头（0表示默认摄像头）
cap = cv2.VideoCapture(0)

# 创建一个画布用来绘制轨迹（与视频同尺寸）
ret, frame = cap.read()
height, width = frame.shape[:2]
trajectory = np.zeros((height, width, 3), dtype=np.uint8)

# 设置红色的HSV阈值（注意红色跨越HSV边界，需两个范围）
lower_red1 = np.array([0, 100, 100])
upper_red1 = np.array([10, 255, 255])
lower_red2 = np.array([160, 100, 100])
upper_red2 = np.array([179, 255, 255])

prev_center = None

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # 翻转镜像，更符合用户习惯
    frame = cv2.flip(frame, 1)

    # 转换颜色空间 BGR -> HSV
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

    # 创建掩膜，检测红色区域
    mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
    mask = cv2.bitwise_or(mask1, mask2)

    # 去噪：先腐蚀再膨胀
    mask = cv2.erode(mask, None, iterations=2)
    mask = cv2.dilate(mask, None, iterations=2)

    # 寻找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 找到最大的轮廓（假设为瓶盖）
    if contours:
        largest_contour = max(contours, key=cv2.contourArea)
        if cv2.contourArea(largest_contour) > 500:  # 设置最小面积
            # 计算轮廓中心
            M = cv2.moments(largest_contour)
            if M["m00"] != 0:
                cX = int(M["m10"] / M["m00"])
                cY = int(M["m01"] / M["m00"])
                center = (cX, cY)

                # 在画布上画轨迹
                if prev_center:
                    cv2.line(trajectory, prev_center, center, (0, 0, 255), 3)
                prev_center = center

                # 在原图中画出中心点和轮廓
                cv2.circle(frame, center, 8, (0, 255, 0), -1)
                cv2.drawContours(frame, [largest_contour], -1, (0, 255, 255), 2)

    # 将轨迹叠加到原图上
    combined = cv2.add(frame, trajectory)

    # 显示图像
    cv2.imshow('Red Cap Tracker', combined)

    # 按q键退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
